What does this PR do ?

Extend policy.megatron_cfg.zero_train_gen_mismatch to FlashInfer MoE experts: per-token NVFP4 (NVFP4 weights × per-token NVFP4 activations) and MXFP8. Train and generation are still bitwise identical (train/gen_kl_error == 0, train/token_mult_prob_error == 1) in a synchronous GRPO loop.

The FlashInfer expert kernel is batch-invariant but not bitwise identical to TE training, so the previous preset rejected it. Megatron-Core's new moe_inference_training_forward makes training use it instead. Each training MoE forward runs generation's FlashInfer kernel, on the same quantization rebuilt from the live parameters on every call. The backward comes from a TE recompute of the MoE layer (a straight-through estimator). For NVFP4, training weights and refit stay BF16; only generation's expert GEMMs (and the matching training forward) run NVFP4.

Builds on #4683 (Nano MXFP8 zero-KL, merged). Shipped: a Nemotron Nano v3 per-token NVFP4 recipe and nightly, and an MXFP8-on-FlashInfer recipe and nightly. Both nightlies stay in disabled.txt until the submodule pins below are on upstream main.

Links

• Builds on (merged): Nano MXFP8 zero-KL, NVIDIA-NeMo/RL#4683
• Merged BF16 zero-KL preset: NVIDIA-NeMo/RL#3208
• Tracking issue: Training-generation mismatch #3743
• Required: FlashInfer train-gen parity and NVFP4 refit: NVIDIA/Megatron-LM#7798
• Already on Megatron-LM main: #8262 (per-token NVFP4 FlashInfer MoE), #7562 (hybrid MoE parity), #8195 (fp64 expert combine), #6076 (deterministic Triton autotune)

What FlashInfer zero-KL needs

Everything the Nano MXFP8 zero-KL preset requires still applies. On top of that:

flashinfer-jit-cache 0.7.0 is a meta-package over per-architecture wheels (flashinfer-jit-cache-sm80 … sm121a) on the explicit flashinfer-cu130 index, which uv only uses for direct dependencies mapped to it. Rather than list and map all of them, the mcore extra drops flashinfer-jit-cache. Per-token NVFP4 doesn't need it (Megatron-Core skips the check for its CuTe DSL JIT kernels), but the FlashInfer MXFP8 path does, so the MXFP8-on-FlashInfer recipe can't run in the default mcore environment yet.

What changed

• nemo_rl/models/megatron/zero_train_gen_mismatch.py:
  • FlashInfer with MXFP8 or NVFP4 experts is allowed when moe_inference_training_forward is on, and _validate_inference_training_forward checks the requirements above.
  • Dedicated generation workers validate with generation_only=True. Their merged config clears the training-only flag, so the "requires moe_inference_training_forward" rule is left to the driver, which sees the full config.
• nemo_rl/models/megatron/setup.py: passes inference_moe_nvfp4_recipe, moe_inference_training_forward and the derived token bound to the model config.
• nemo_rl/models/generation/megatron/config.py: merged_inference_megatron_cfg clears moe_inference_training_forward for generation.
• nemo_rl/models/policy/__init__.py: the three new MegatronConfig keys; te_precision_config_file accepts null.
• nemo_rl/models/policy/workers/megatron_policy_worker.py: passes generation_only=is_refit_destination.
• pyproject.toml / uv.lock: FlashInfer 0.7.0.post1 in the mcore extra, without flashinfer-jit-cache; the lock also picks up the Bridge bump (nccl-extensions from Megatron-LM, Pillow 13).
• Recipes and nightlies: grpo-nanov3-30BA3B-2n4g-megatron_generation-noncolocated-zero-kl-mxfp8-flashinfer (on the Nano MXFP8 zero-KL recipe) and ...-zero-kl-nvfp4-flashinfer (directly on grpo_math_1B.yaml), both in tests/test_suites/disabled.txt for now.
• Unit tests for the new validator rules, the generation-only path and the token bound.

Validation

• Unit: tests/unit/models/megatron/test_zero_train_gen_mismatch.py, tests/unit/models/megatron/test_megatron_setup.py.
• Functional (GB200, 2 nodes × 4 GPUs, 5 sync GRPO steps, Nemotron Nano v3 per-token NVFP4 nightly, at the pins above): job 1309408 passed (W&B). train/gen_kl_error was exactly 0, and train/token_mult_prob_error, train/probs_ratio_min and train/probs_ratio_max were exactly 1 on all 5 steps.
• The MXFP8-on-FlashInfer nightly has not run at these pins; it needs flashinfer-jit-cache (above).